# Author: 左志华
# Example: python ampd.py --file demo\*.csv -n 1
# Input: 采用命令行输入。
#   1. `--file`指定时序文件，第一行为表头，会被忽视，第一列为时间，第二列为时序信号；
#   2. `--save`指定是否保存图像，默认不保存；
#   3. `--ncols`指定时序文件列数，默认为1，即时间和信号两列。
# Output: 查询峰谷值，并获取幅值。
# Function: 对气垫船时序进行绘制和获取物理量幅值均值。
#   1. 读取气垫船时序数据;
#   2. 获取峰谷值，计算均值，获取幅值;
#   3. 绘制图像，保存图像。
# Note: 算法存在局限性，仅适用于周期性时序。
import numpy as np
import matplotlib.pyplot as plt
import argparse
import scipy.signal as singal
import scienceplots

parser = argparse.ArgumentParser()
parser.add_argument("-f", "--file", help="Time series filename", required=True)
parser.add_argument("-s", "--save", help="Save figure", action="store_true")
parser.add_argument("-n", "--ncols", type=int, help="Number of columns", default=1)
args = parser.parse_args()

print("FILTER v1.0")
print("Copyright (c) 2023, ZUO Zhihua (E-mail: zuo.zhihua@qq.com)\n")

print("Input: ", args.file)
print("Save figure: ", args.save, "\n")

# 读取数据并提取稳定阶段时间序列
data = np.loadtxt(args.file, delimiter=",", skiprows=1)
stable_time = data[:, 0]
stable_singal = data[:, args.ncols]


def AMPD(data):
    """
    Automatic multiscale-based peak detection (AMPD) algorithm

    AMPD 峰值算法

    Parameters
    ---
    data : array_like
        输入的一维数据

    Returns
    ---
    p_data : ndarray
        波峰所在索引值的列表
    """
    p_data = np.zeros_like(data, dtype=np.int32)
    N = data.shape[0]
    L = N // 2

    # 计算每个窗口的行和，并存储到arr_rowsum数组中
    arr_rowsum = np.zeros(L, dtype=np.int32)
    for k in range(1, L):
        arr_rowsum[k] = np.sum(
            (data[k:-k] > data[k * 2 :]) & (data[k:-k] > data[: -k * 2])
        )

    # 找到最大的行和对应的窗口长度
    max_index = np.argmax(arr_rowsum)

    # 根据窗口长度计算波峰
    for k in range(1, max_index + 1):
        p_data[k:-k][
            (data[k:-k] > data[k * 2 :]) & (data[k:-k] > data[: -k * 2])
        ] += 1  # 花式索引

    return np.where(p_data == max_index)[0]


# 查询峰谷值，并取均值
peaks = AMPD(stable_singal)
troughs = AMPD(-stable_singal)

peak_mean = np.mean(stable_singal[peaks[2:-1]])
trough_mean = np.mean(stable_singal[troughs[2:-1]])

print("peaks' mean = ", peak_mean)
print("troughs' mean = ", trough_mean)
print("**amplitude** = ", (peak_mean - trough_mean) / 2)

# 绘制滤波后的结果
with plt.style.context(["science", "no-latex"]):
    fig, ax = plt.subplots(figsize=(3.0, 1.7))  # 设置图形大小
    plt.plot(stable_time, stable_singal)
    plt.xlabel("Time (s)")
    plt.ylabel("Signal")

    plt.plot(stable_time[peaks], stable_singal[peaks], "x", label="peaks")
    plt.plot(stable_time[troughs], stable_singal[troughs], "o", label="troughs")

    # 绘制平均水平线
    plt.axhline(peak_mean, color="r", linestyle="--", label="peaks' mean")
    plt.axhline(trough_mean, color="g", linestyle="--", label="troughs' mean")

    # 添加图例和标题
    plt.legend(facecolor="white", framealpha=0.75, frameon=True, loc="best")
    if args.save:
        plt.savefig("peaks_and_troughs.svg")
        plt.close()
    else:
        plt.show()
